const fs = require('fs').promises;
const path = require('path');

/**
 * 编辑类型枚举
 */
const EditType = {
    UNDEFINED: 0,
    NO_FILE_READ: 1,
    READ_FILE_TOOL: 2,
    ALTER_FILE_TOOL: 3,
    FILE_MENTION: 4
};

/**
 * 上下文管理器 - 负责管理和优化聊天上下文
 */
class ContextManager {
    constructor() {
        // 上下文历史更新映射
        // 格式: { outerIndex => [EditType, { innerIndex => [[timestamp, updateType, update, metadata], ...] }] }
        this.contextHistoryUpdates = new Map();
    }

    /**
     * 初始化上下文历史
     * @param {string} taskDirectory - 任务目录路径
     * @returns {Promise<void>}
     */
    async initializeContextHistory(taskDirectory) {
        this.contextHistoryUpdates = await this.getSavedContextHistory(taskDirectory);
    }

    /**
     * 从磁盘获取保存的上下文历史
     * @param {string} taskDirectory - 任务目录路径
     * @returns {Promise<Map<number, [number, Map<number, Array>]>>} - 上下文历史更新映射
     */
    async getSavedContextHistory(taskDirectory) {
        try {
            const filePath = path.join(taskDirectory, 'context_history.json');
            try {
                // 检查文件是否存在
                await fs.access(filePath);
                const data = await fs.readFile(filePath, 'utf8');
                const serializedUpdates = JSON.parse(data);
                
                // 重建映射结构
                return new Map(
                    serializedUpdates.map(([messageIndex, [numberValue, innerMapArray]]) => [
                        messageIndex,
                        [numberValue, new Map(innerMapArray)]
                    ])
                );
            } catch (error) {
                if (error.code === 'ENOENT') {
                    return new Map(); // 文件不存在，返回空Map
                }
                throw error;
            }
        } catch (error) {
            console.error('加载上下文历史失败:', error);
            return new Map();
        }
    }

    /**
     * 保存上下文历史到磁盘
     * @param {string} taskDirectory - 任务目录路径
     * @returns {Promise<void>}
     */
    async saveContextHistory(taskDirectory) {
        try {
            // 序列化映射结构
            const serializedUpdates = Array.from(this.contextHistoryUpdates.entries()).map(
                ([messageIndex, [numberValue, innerMap]]) => [messageIndex, [numberValue, Array.from(innerMap.entries())]]
            );
            
            // 写入文件
            await fs.writeFile(
                path.join(taskDirectory, 'context_history.json'),
                JSON.stringify(serializedUpdates),
                'utf8'
            );
        } catch (error) {
            console.error('保存上下文历史失败:', error);
        }
    }

    /**
     * 获取更新的上下文消息和元数据
     * @param {Array} apiConversationHistory - API对话历史
     * @param {Array} uiMessages - UI消息
     * @param {any} api - API接口
     * @param {Array} conversationHistoryDeletedRange - 已删除的对话历史范围
     * @param {number} previousApiReqIndex - 上一个API请求索引
     * @param {string} taskDirectory - 任务目录路径
     * @returns {Promise<Object>} - 更新的上下文信息
     */
    async getNewContextMessagesAndMetadata(
        apiConversationHistory,
        uiMessages,
        api,
        conversationHistoryDeletedRange,
        previousApiReqIndex,
        taskDirectory
    ) {
        let updatedConversationHistoryDeletedRange = false;
        let anyContextUpdates = false;

        // 检查上下文窗口大小和可能的截断需求
        if (previousApiReqIndex >= 0) {
            const previousRequest = uiMessages[previousApiReqIndex];
            if (previousRequest && previousRequest.text) {
                const timestamp = previousRequest.ts;
                const requestInfo = JSON.parse(previousRequest.text);
                const totalTokens = (requestInfo.tokensIn || 0) + (requestInfo.tokensOut || 0);
                
                // 如果令牌总数接近上下文窗口限制，可能需要截断
                const maxAllowedSize = this.getContextWindowInfo(api);
                
                if (totalTokens >= maxAllowedSize * 0.8) { // 如果已经使用了80%的上下文窗口
                    // 首先尝试通过优化来减少令牌使用
                    const [hasContextUpdates, uniqueFileReadIndices] = this.applyContextOptimizations(
                        apiConversationHistory,
                        conversationHistoryDeletedRange ? conversationHistoryDeletedRange[1] + 1 : 2,
                        timestamp
                    );
                    
                    anyContextUpdates = hasContextUpdates;
                    
                    // 检查优化是否足够，如果不够就进行截断
                    let needToTruncate = true;
                    if (anyContextUpdates) {
                        const charactersSavedPercentage = this.calculateContextOptimizationMetrics(
                            apiConversationHistory,
                            conversationHistoryDeletedRange,
                            uniqueFileReadIndices
                        );
                        
                        if (charactersSavedPercentage >= 0.3) { // 如果节省了30%以上的字符，不需要截断
                            needToTruncate = false;
                        }
                    }
                    
                    if (needToTruncate) {
                        // 添加截断通知
                        const hasAppliedNotice = this.applyStandardContextTruncationNoticeChange(timestamp);
                        anyContextUpdates = anyContextUpdates || hasAppliedNotice;
                        
                        // 获取下一个截断范围
                        conversationHistoryDeletedRange = this.getNextTruncationRange(
                            apiConversationHistory,
                            conversationHistoryDeletedRange,
                            totalTokens / 2 > maxAllowedSize ? 'quarter' : 'half'
                        );
                        
                        updatedConversationHistoryDeletedRange = true;
                    }
                    
                    // 如果有任何上下文更新，保存到磁盘
                    if (anyContextUpdates) {
                        await this.saveContextHistory(taskDirectory);
                    }
                }
            }
        }
        
        // 获取并应用上下文更改的截断消息
        const truncatedConversationHistory = this.getAndAlterTruncatedMessages(
            apiConversationHistory,
            conversationHistoryDeletedRange
        );
        
        return {
            conversationHistoryDeletedRange,
            updatedConversationHistoryDeletedRange,
            truncatedConversationHistory
        };
    }
    
    /**
     * 获取上下文窗口信息
     * @param {any} api - API接口
     * @returns {number} - 最大允许的上下文大小
     */
    getContextWindowInfo(api) {
        // 根据API选择适当的上下文窗口大小
        // 这里提供一个默认值，需要根据实际使用的模型来调整
        return 4000; // 默认值
    }

    /**
     * 获取下一个截断范围
     * @param {Array} apiMessages - API消息
     * @param {Array} currentDeletedRange - 当前删除范围
     * @param {string} keep - 保留策略 ('half' 或 'quarter')
     * @returns {Array} - 新的删除范围
     */
    getNextTruncationRange(apiMessages, currentDeletedRange, keep) {
        // 我们始终保留第一个用户-助手对话，从那里开始截断偶数数量的消息
        const rangeStartIndex = 2; // 索引0和1被保留
        const startOfRest = currentDeletedRange ? currentDeletedRange[1] + 1 : 2; // 包含起始索引
        
        // 计算要保留的消息数量
        const remainingMessages = apiMessages.length - startOfRest;
        const keepCount = keep === 'quarter' ? Math.floor(remainingMessages / 4) : Math.floor(remainingMessages / 2);
        
        // 确保我们保留偶数数量的消息（用户-助手对）
        const adjustedKeepCount = keepCount - (keepCount % 2);
        
        // 计算新的截断范围
        const newRangeStart = rangeStartIndex;
        const newRangeEnd = apiMessages.length - adjustedKeepCount - 1;
        
        return [newRangeStart, newRangeEnd];
    }

    /**
     * 获取并修改截断的消息
     * @param {Array} messages - 消息数组
     * @param {Array} deletedRange - 删除范围
     * @returns {Array} - 修改后的消息数组
     */
    getAndAlterTruncatedMessages(messages, deletedRange) {
        const truncatedMessages = this.getTruncatedMessages(messages, deletedRange);
        return this.applyContextHistoryUpdates(truncatedMessages, 0);
    }

    /**
     * 获取截断的消息
     * @param {Array} messages - 消息数组
     * @param {Array} deletedRange - 删除范围
     * @returns {Array} - 截断后的消息数组
     */
    getTruncatedMessages(messages, deletedRange) {
        if (!deletedRange) {
            return [...messages];
        }
        
        const [start, end] = deletedRange;
        return [...messages.slice(0, start), ...messages.slice(end + 1)];
    }

    /**
     * 应用上下文历史更新
     * @param {Array} messages - 消息数组
     * @param {number} startFromIndex - 起始索引
     * @returns {Array} - 更新后的消息数组
     */
    applyContextHistoryUpdates(messages, startFromIndex) {
        const result = [...messages];
        
        // 遍历上下文历史更新映射
        for (const [messageIndex, [editType, innerMap]] of this.contextHistoryUpdates.entries()) {
            if (messageIndex < startFromIndex || messageIndex >= result.length) {
                continue;
            }
            
            const message = result[messageIndex];
            
            // 根据消息类型处理
            if (message.role === 'assistant' && typeof message.content === 'object') {
                // 处理助手消息的内容块
                if (Array.isArray(message.content)) {
                    for (const [blockIndex, updates] of innerMap.entries()) {
                        if (blockIndex < message.content.length) {
                            const block = message.content[blockIndex];
                            
                            if (block.type === 'text' && updates.length > 0) {
                                // 获取最后的更新（最新的）
                                const lastUpdate = updates[updates.length - 1];
                                const [, , update] = lastUpdate;
                                
                                // 应用更新
                                block.text = update[0] || '';
                            }
                        }
                    }
                }
            } else if (typeof message.content === 'string') {
                // 处理字符串内容
                // 简单处理，映射索引0作为整个消息内容的更新
                const updates = innerMap.get(0);
                if (updates && updates.length > 0) {
                    const lastUpdate = updates[updates.length - 1];
                    const [, , update] = lastUpdate;
                    
                    // 应用更新
                    message.content = update[0] || '';
                }
            }
        }
        
        return result;
    }

    /**
     * 在指定时间戳截断上下文历史
     * @param {number} timestamp - 时间戳
     * @param {string} taskDirectory - 任务目录路径
     * @returns {Promise<void>}
     */
    async truncateContextHistory(timestamp, taskDirectory) {
        this.truncateContextHistoryAtTimestamp(this.contextHistoryUpdates, timestamp);
        await this.saveContextHistory(taskDirectory);
    }

    /**
     * 根据时间戳截断上下文历史
     * @param {Map} contextHistory - 上下文历史映射
     * @param {number} timestamp - 时间戳
     */
    truncateContextHistoryAtTimestamp(contextHistory, timestamp) {
        for (const [messageIndex, [editType, innerMap]] of contextHistory.entries()) {
            for (const [blockIndex, updates] of innerMap.entries()) {
                // 找到大于给定时间戳的更新的索引
                const cutIndex = updates.findIndex(update => update[0] > timestamp);
                
                if (cutIndex !== -1) {
                    // 删除所有更新的索引
                    updates.splice(cutIndex);
                    
                    // 如果没有更新，删除块条目
                    if (updates.length === 0) {
                        innerMap.delete(blockIndex);
                    }
                }
            }
            
            // 如果内部映射为空，删除消息条目
            if (innerMap.size === 0) {
                contextHistory.delete(messageIndex);
            }
        }
    }

    /**
     * 应用上下文优化
     * @param {Array} apiMessages - API消息
     * @param {number} startFromIndex - 起始索引
     * @param {number} timestamp - 时间戳
     * @returns {Array} - [是否有任何上下文更新, 唯一文件读取索引集合]
     */
    applyContextOptimizations(apiMessages, startFromIndex, timestamp) {
        // 查找并潜在地保存文件读取上下文历史更新
        return this.findAndPotentiallySaveFileReadContextHistoryUpdates(
            apiMessages,
            startFromIndex,
            timestamp
        );
    }

    /**
     * 应用标准上下文截断通知更改
     * @param {number} timestamp - 时间戳
     * @returns {boolean} - 是否应用了更改
     */
    applyStandardContextTruncationNoticeChange(timestamp) {
        // 将通知添加到第一个助手消息
        const firstAssistantMessageIndex = 1; // 假设第一个消息是用户，第二个是助手
        let messageEntry = this.contextHistoryUpdates.get(firstAssistantMessageIndex);
        
        if (!messageEntry) {
            messageEntry = [EditType.NO_FILE_READ, new Map()];
            this.contextHistoryUpdates.set(firstAssistantMessageIndex, messageEntry);
        }
        
        let blockUpdates = messageEntry[1].get(0);
        if (!blockUpdates) {
            blockUpdates = [];
            messageEntry[1].set(0, blockUpdates);
        }
        
        // 添加截断通知
        blockUpdates.push([
            timestamp,
            'text',
            ['[注意] 为了节省空间，部分历史对话已被移除。']
        ]);
        
        return true;
    }

    /**
     * 计算上下文优化指标
     * @param {Array} apiMessages - API消息
     * @param {Array} conversationHistoryDeletedRange - 对话历史删除范围
     * @param {Set} uniqueFileReadIndices - 唯一文件读取索引集合
     * @returns {number} - 节省的字符百分比
     */
    calculateContextOptimizationMetrics(apiMessages, conversationHistoryDeletedRange, uniqueFileReadIndices) {
        // 简化版本：假设优化节省了10%的字符
        return 0.1;
    }

    /**
     * 查找并潜在地保存文件读取上下文历史更新
     * @param {Array} apiMessages - API消息
     * @param {number} startFromIndex - 起始索引
     * @param {number} timestamp - 时间戳
     * @returns {Array} - [是否有任何上下文更新, 唯一文件读取索引集合]
     */
    findAndPotentiallySaveFileReadContextHistoryUpdates(apiMessages, startFromIndex, timestamp) {
        // 简化版本：返回没有更新
        return [false, new Set()];
    }
}

module.exports = ContextManager; 